import matplotlib.pyplot as plt
from PIL import Image


def sample_img(input_img, output_img, epoch):
    '''
    每过一个epoch就让我康康结果
    :param input_img: 输入图片
    :param output_img: 输出图片
    :param epoch: 当前训练轮数
    :return:
    '''
    # 声明一个图框 ----------------------------------------------

    fig = plt.figure()

    # 往图框中插入一个子图， 1行2列的大小，放在第1个位置----
    axl = fig.add_subplot(121)
    axl.set_title('input img')      # 设置子图标题
    axl.axis('off')                 # 子图是否开启坐标 on/off
    axl.imshow(input_img)           # 让子图展示ta
    # ----------------------------------------------

    # 往图框中插入一个子图， 1行2列的大小，放在第2个位置----
    axl = fig.add_subplot(122)
    axl.set_title('output img')     # 设置子图标题
    axl.axis('off')                 # 子图是否开启坐标 on/off
    axl.imshow(output_img)          # 让子图展示ta
    # ----------------------------------------------

    # ----------------------------------------------------------

    # 保存图框
    fig.savefig("result_{0}.png".format(epoch))


if __name__ == '__main__':

    # 随便整点模拟数据 -------------------------
    input_img = Image.open('input_img.jpg')
    output_img = Image.open('output_img.jpg')
    epoch = 5
    # ---------------------------------------

    # 保存训练结果
    sample_img(input_img, output_img, epoch)
